import onnxruntime    
import onnx
import numpy as np

# class VocOnnxRuntime():
#     def __init__(self, onnx_model, providers=[]):
        
#         sess_options = onnxruntime.SessionOptions()
#         sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
#         sess_options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
#         sess_options.intra_op_num_threads = cpu_threads
#         self.model = onnxruntime.InferenceSession(onnx_modle, providers=providers, sess_options=sess_options)
#         self.warming_up()
    
#     def forward(self, mel):
#         input_dict = {"mels":mel}
#         output = self.model.run(None, input_dict)
#         wave = output[0]
#         print(output.shape)
#         print(wave.shape)
#         return wave
    
#     def inference(self, mel):
#         return self.forward(mel)
    
#     def warming_up(self):
#         mels = np.random.randint(1, 32, (1, 100, 80))/1000.0
#         mels = mels.astype(np.float32)
#         _ = self.infernece(mels)



# 指定你的 ONNX 模型文件路径
onnx_model_path = "hf_new.onnx"

# 加载 ONNX 模型
onnx_model = onnx.load(onnx_model_path)

# 创建 ONNX Runtime 推理会话
session = onnxruntime.InferenceSession(onnx_model_path)

# 检查输入节点和输出节点的名称
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

print(input_name)
print(output_name)

# 创建一个示例输入数据（这里假设输入是一个大小为 (1, input_dim) 的二维数组）
mels = np.random.randint(1, 32, (3, 1000, 80))/1000.0
mels = mels.astype(np.float32)
print("mels shape:{}".format(mels.shape))
# input_data = np.random.randn(1, input_dim).astype(np.float32)

# 进行推理
output_data = session.run([output_name], {input_name: mels})

print(len(output_data))
print(output_data[0].shape)

# 输出推理结果
print("Inference result:", output_data)

    

# if __name__ == '__main__':
#     onnxruntime = VocOnnxRuntime('/root/zhanggj/hifi-gan/hf_new.onnx')
    # onnxruntime

    # 加载模型
    # model = onnx.load('/root/zhanggj/hifi-gan/hf_new.onnx')
    # 检查模型格式是否完整及正确
    # onnx.checker.check_model(model)
    # 获取输出层，包含层名称、维度信息
    # output = model.graph.input
    # print(output)
